#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
使用 backtrader 框架读取股票数据并打印 BOLL 指标
"""

import os
import json
import pandas as pd
import backtrader as bt
from pathlib import Path
from loguru import logger
from dotenv import load_dotenv
from datetime import datetime

# 加载环境变量
load_dotenv()

def data_in_upchannel(data, n=6, threshold=4):
    return len([i for i in range(n) if data[-i]>=data[-i-1]]) >= threshold


class BollDataStrategy(bt.Strategy):
    """打印 BOLL 指标数据的策略"""
    params = (
        ('target_trade_date', "2025-07-25"),
    )
    def __init__(self):
        if type(self.params.target_trade_date) == str:
            self.params.target_trade_date = datetime.strptime(self.params.target_trade_date, '%Y-%m-%d').date()
    
    def next(self):
        current_date = self.datas[0].datetime.date(0)
        if current_date != self.params.target_trade_date:
            return
        result={}
        logger.info(f"当前日期: {current_date} 是目标日期: {self.params.target_trade_date}")
        result['target_trade_date']=str(current_date)
        result['stock_meet_all']=[]
        
        for data in self.datas:
            logger.debug(f"{data._name} 当前日期: {data.datetime.date(0)}, BOLL: {data.boll[0]:.2f}, UB: {data.ub[0]:.2f}, LB: {data.lb[0]:.2f}")
            if all(data_in_upchannel(indicator) for indicator in [data.boll, data.ub, data.lb]):
                result['stock_meet_all'].append(data._name)
        logger.info(f"符合条件的股票数量及详情: {len(result['stock_meet_all'])} {result['stock_meet_all']}")

        result['count']=len(result['stock_meet_all'])
        
        # 把result写入文件，路径是{project_root}/data/result1.json。使用pathlib
        project_root = Path(__file__).parent.parent
        result_file = project_root / 'data' / 'result1.json'
        result_file.parent.mkdir(exist_ok=True)  # 确保data目录存在
        with open(result_file, 'w', encoding='utf-8') as f:
            json.dump(result, f, ensure_ascii=False, indent=2)


class BollDataFeed(bt.feeds.PandasData):
    """自定义数据源，包含 BOLL 指标"""
    
    # 定义额外的列
    lines = ('boll', 'ub', 'lb')
    
    params = (
        ('boll', 'BOLL-M.BOLL'),
        ('ub', 'BOLL-M.UB'),
        ('lb', 'BOLL-M.LB'),
    )

def load_stock_data(file_path):
    """加载股票数据文件"""
    try:
        # 读取文件，跳过前两行（标题行），使用制表符分隔
        df = pd.read_csv(file_path, encoding='gbk', skiprows=2, sep='\t')
        
        # 重命名列名，去除空格
        df.columns = df.columns.str.strip()
        # logger.debug(f"列名: {df.columns}")
        
        # 过滤掉非日期行（包含非数字字符的行）
        df = df[df['时间'].str.match(r'^\s*\d{4}/\d{2}/\d{2}\s*$', na=False)]
        
        # 处理日期列，去除前后空格
        df['时间'] = df['时间'].str.strip()
        df['时间'] = pd.to_datetime(df['时间'], format='%Y/%m/%d')
        df.set_index('时间', inplace=True)
        
        # 确保 BOLL 指标列存在
        required_columns = ['BOLL-M.BOLL', 'BOLL-M.UB', 'BOLL-M.LB']
        for col in required_columns:
            if col not in df.columns:
                logger.warning(f"列 {col} 不存在，使用默认值 0")
                df[col] = 0.0
        
        # 处理空值
        df[required_columns] = df[required_columns].fillna(0.0)
        
        # 确保价格列存在
        if '收盘' not in df.columns:
            logger.error("缺少收盘价列")
            return None
        
        # 转换数值列为数值类型
        numeric_columns = ['开盘', '最高', '最低', '收盘', '成交量'] + required_columns
        for col in numeric_columns:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors='coerce').fillna(0.0)
            
        return df
        
    except Exception as e:
        logger.error(f"加载文件 {file_path} 时出错: {e}")
        return None

def main():
    """主函数"""
    logger.info("开始处理股票数据...")
    
    # 从环境变量获取股票数据路径
    stock_path = os.getenv('GUANGFA_STOCK_PATH')
    if not stock_path:
        logger.error("未找到 GUANGFA_STOCK_PATH 环境变量")
        return
    
    stock_path = Path(stock_path)
    if not stock_path.exists():
        logger.error(f"路径不存在: {stock_path}")
        return
    
    logger.info(f"股票数据路径: {stock_path}")
    
    # 查找所有 txt 文件
    txt_files = list(stock_path.glob("*.txt"))
    if not txt_files:
        logger.warning(f"在 {stock_path} 中未找到任何 txt 文件")
        return
    
    logger.info(f"找到 {len(txt_files)} 个股票数据文件")
    
    # 处理所有文件
    total_files = len(txt_files)
    processed_files = 0
    failed_files = 0

    # 创建 backtrader 引擎
    cerebro = bt.Cerebro()

    for i, file_path in enumerate(txt_files, 1):
        stock_code = file_path.stem  # 获取文件名（不含扩展名）
        logger.info(f"处理股票 {i}/{total_files}: {stock_code}")
        
        # 加载数据
        df = load_stock_data(file_path)
        if df is None:
            logger.error(f"数据加载失败: {stock_code}")
            failed_files += 1
            continue
        
        logger.info(f"数据加载成功，共 {len(df)} 条记录")
        logger.info(f"数据时间范围: {df.index.min()} 到 {df.index.max()}")
        
                
        # 添加数据源
        data_feed = BollDataFeed(
            dataname=df,
            datetime=None,  # 使用索引作为日期
            open='开盘',
            high='最高',
            low='最低',
            close='收盘',
            volume='成交量',
            openinterest=-1,  # 不使用持仓量
            boll='BOLL-M.BOLL',
            ub='BOLL-M.UB',
            lb='BOLL-M.LB'
        )
        
        cerebro.adddata(data_feed, name=stock_code)
        processed_files += 1

        
    # 添加策略
    cerebro.addstrategy(BollDataStrategy, target_trade_date=datetime.strptime(df.index.max().strftime('%Y-%m-%d'), '%Y-%m-%d').date())
        
        
    # 运行回测
    logger.info(f"开始运行 backtrader 回测...")
    cerebro.run()
        
    
    logger.info(f"所有文件处理完成！成功处理: {processed_files} 个文件，失败: {failed_files} 个文件")
    logger.info("BOLL 指标数据打印完成")

if __name__ == "__main__":
    main()
